import numpy as np
import json

class TimeAxis():
    def __init__(self, other=None):
        self.times, self.nfr_list = [], []
        self.clipfr2time, self.time2clipfr = {}, {}
    
    def get_shape(self):
        return self.times.shape if self.times is np.array else len(self.times)

    def is_contiguous(self):
        '''
            判断时间轴是否为连续
            一般使用interpolates.interp_by_xxx来完成此操作
        '''
        prev_t = -1
        for t in self.times:
            if t - prev_t > 1:
                return False
            prev_t = t
        return True

    def next_contiguous_clipfr(self, i_t):
        '''
            获取下一个时间点的clip和fr编号
            return next_clip_idx, next_fr_idx
        '''
        t = self.times[i_t]
        clip, fr = self.clipfr2time[t]
        nfr = self.nfr_list[clip]
        if fr < nfr - 1:
            return clip, fr+1
        else:
            return clip+1, 0

    def slice_by(self, start=0, stop=-1):
        self.times = self.times[start:stop]

    def scale_data_by(data:np.array, interv):
        pass


def create_timeaxis_from(f_info, clipfr_list:np.array, fps=1):
    '''
        建立 (clip,fr) -> absolute_time 双向字典
        param.f_info: 包含每隔视频片段全长列表的属性文件
        param.clipfr_list: [[clip_idx, fr_idx]], Nx2
    '''
    # 获取每个clip的总长
    nfr_list = []
    with open(f_info, 'r') as h:
        df = json.load(h)
        nfr_list = [int(nfr) for nfr in df['nfr_list']]
    # 计算绝对时间轴
    prev_clip, prev_fr, prev_nfr = None, None, None
    times = [] # 绝对时间轴
    curr_time = 0 # 当前绝对时间
    for clipfr in clipfr_list:
        clip, fr = clipfr # 获取当前视频片段和帧编号
        nfr = nfr_list[clip] # 获取当前视频片段总长
        if prev_fr is not None: # 若不是第一帧，则开始检查丢帧
            if clip != prev_clip and prev_fr < prev_nfr - fps: # 若在相邻视频片段之间丢帧
                lost_sec = (prev_nfr - fps - prev_fr) / fps # 计算所丢失的时长
                curr_time += int(lost_sec) # 跳过所丢时长
            elif fr - prev_fr > fps: # 若在视频内丢帧
                lost_sec  = (fr - fps - prev_fr) / fps
                curr_time += int(lost_sec)
        times.append(curr_time)
        # 更新所有prev数据
        curr_time += 1
        prev_clip, prev_fr, prev_nfr = clip, fr, nfr
    # 创建TimeAxis
    taxis = TimeAxis()
    # 保存time axis
    taxis.times = times
    taxis.nfr_list = nfr_list
    # 组建双向字典
    for clipfr,t in zip(clipfr_list, times):
        taxis.clipfr2time[tuple(clipfr)]=t
        taxis.time2clipfr[t] = clipfr
    return taxis


def interp_by_lastcopy(taxis:TimeAxis, values, interv=1):
    '''
        使用最近数据插值
        times, values 分别为原始时间轴和数据轴
        interv 为单位时间间隔
        return 插值后的数据轴
    '''
    if len(taxis.times) != len(values):
        print('[interp_by_lastcopy.wargning] mis-aligned times and values.')
        return None, None
    prev_t, prev_v = -1, None
    res_val = []
    clipfr_list = []
    n_clip = len(taxis.nfr_list)
    for t, v in zip(taxis.times, values):
        diff = int(t - prev_t)
        if prev_t > -1 and diff > interv: # 若跳帧，需插值
            interps = [prev_v for i in range(diff)]
            res_val.extend(interps)
        else: # 连续帧，直接复制
            res_val.append(v)
        # 刷新prev数值
        prev_t,  prev_v = t, v
    # 对最后视频末尾插帧
    t_last = taxis.times[-1]
    clip_last, fr_last = taxis.time2clipfr[t_last]
    diff = taxis.nfr_list[-1] - fr_last
    v_last = values[-1]
    res_val.extend(v_last for i in range(diff))
    return np.mat(res_val)


def get_contiguous_timeaxis(taxis:TimeAxis):
    '''
        根据输入的time axis创建一个连续的time axis
    '''
    times = []
    offset = 0
    clip_list = []
    fr_list = []
    for clip, nfr in enumerate(taxis.nfr_list): # 根据每个clip的nfr
        times.extend(list(range(offset, offset+nfr))) # 构造连续时间轴
        clip_list.extend([clip] * nfr) # 保存连续clip轴
        fr_list.extend(list(range(nfr))) # 保存连续fr轴
        offset += nfr # 更新绝对时间轴的时间偏移量
    # 构造新的时间轴对象
    new_taxis = TimeAxis()
    new_taxis.times = times
    new_taxis.nfr_list = taxis.nfr_list
    # 保存新轴的双向字典
    for t, clip, fr in zip(times, clip_list, fr_list):
        new_taxis.clipfr2time[(clip, fr)] = t
        new_taxis.time2clipfr[t] = (clip,fr)
    return new_taxis


def load_tracing_data(f_loc, f_info):
    '''
        载入追踪数据，矩阵格式
        clip    xx  fr  xx  body    x   y   w   h   head    x   y   w   h
        return time_axis, body_loc, head_locc
    '''
    lines = []
    with open(f_loc,'r') as h: # 读取locations.txt每行信息
        lines = [line.rstrip('\r\t\n ') for line in h.readlines()]
    mat = []
    ttypes = ['clip', 'fr', 'body', 'head']
    for line in lines:
        terms = line.split('\t')
        curr_data = []
        for term in terms: # 分解当前帧的数据条目
            if term in ttypes:
                continue
            else:
                curr_data.append(int(term)) # 一帧数据放一行
        mat.append(curr_data)
    mat = np.array(mat) # Nx10
    clipfr_list, bodyhead_loc = mat[:, :2], mat[:, 2:] # 按内容分开
    # 创建时间轴
    taxis = create_timeaxis_from(f_info, clipfr_list)
    return taxis, bodyhead_loc


def locate_1st_different_pt(times1, times2):
    stop = (-1, -1)
    for t1, t2 in zip(times1, times2):
        if t1 != t2:
            stop=(t1, t2)
            break
    stop = (t1, t2)
    return stop